#!/bin/bash
set -o errexit

#
# Identify the version of this USB key (if it is indeed a USB key).
#
# We do this by sniffing fixed offset within the MBR. If we have a (legacy)
# grub MBR, then we can look at offset 0x3e for COMPAT_VERSION_MAJOR and
# COMPAT_VERSION_MINOR (and we'll presume 3.2 as a minimum).
#
# If we're talking about a loader-based key, we'll peek at 0xfa AKA
# STAGE1_MBR_VERSION for format_image's IMAGE_MAJOR, which we expect to be 2.
#
# Unfortunately there's no standard way to find a version for other MBRs such as
# grub2's. In these cases we'll end up with a potentially random version here,
# so this key should not be trusted as ours until mounted and the path
# .joyusbkey is found.
#
function usb_key_version()
{
        local readonly devpath=$1
        local readonly mbr_sig_offset=0x1fe
        local readonly mbr_grub_offset=0x3e
        local readonly mbr_stage1_offset=0xfa
        local readonly mbr_grub_version=0203
        local readonly mbr_sig=aa55

        sig=$(echo $(/usr/bin/od -t x2 \
            -j $mbr_sig_offset -A n -N 2 $devpath) )

        if [[ "$sig" != $mbr_sig ]]; then
                echo "unknown"
                return
        fi

        grub_val=$(echo $(/usr/bin/od -t x2 \
            -j $mbr_grub_offset -A n -N 2 $devpath) )
        loader_major=$(echo $(/usr/bin/od -t x1 \
            -j $mbr_stage1_offset -A n -N 1 $devpath) )

        if [[ "$grub_val" = $mbr_grub_version ]]; then
                echo "1"
                return
        fi

        echo $(( 0x$loader_major ))
}

#
# Mount the usbkey at the standard mount location (or whatever is specified).
#
function mount_usb_key()
{
        local mnt=$1

        if [[ -z "$mnt" ]]; then
                mnt=/mnt/$(svcprop -p "joyentfs/usb_mountpoint" \
                    "svc:/system/filesystem/smartdc:default")
        fi

        if [[ -f "$mnt/.joyliveusb" ]]; then
                echo $mnt
                return 0
        fi

        if ! mkdir -p $mnt; then
                echo "failed to mkdir $mnt" >&2
                return 1
        fi

#       readonly alldisks=$(/usr/bin/disklist -a)
        readonly alldisks=$(/usr/bin/disklist -r)

        for disk in $alldisks; do
                version=$(usb_key_version "/dev/dsk/${disk}p0")

                case $version in
                1) devpath="/dev/dsk/${disk}p1" ;;
                2) devpath="/dev/dsk/${disk}s2" ;;
                *) continue ;;
                esac

                fstyp="$(/usr/sbin/fstyp $devpath 2>/dev/null)"

                if [[ "$fstyp" != "pcfs" ]]; then
                        continue
                fi

                /usr/sbin/mount -F pcfs -o foldcase,noatime $devpath $mnt \
                    2>/dev/null

                if [[ $? -ne 0 ]]; then
                        continue
                fi

#               if [[ -f $mnt/.joyliveusb ]]; then
                        echo $mnt
                        return 0
#               fi

                if ! /usr/sbin/umount $mnt; then
                        echo "Failed to unmount $mnt" >&2
                        return 1
                fi
        done

        echo "Couldn't find USB key" >&2
        return 1
}

function unmount_usb_key()
{
        local mnt=$1

        if [[ -z "$mnt" ]]; then
                mnt=/mnt/$(svcprop -p "joyentfs/usb_mountpoint" \
                    "svc:/system/filesystem/smartdc:default")
        fi

        typ=$(awk -v "mnt=$mnt" '$2 == mnt { print $3 }' /etc/mnttab)

        [[ -z $typ ]] && return 0

#       if [[ ! -f "$mnt/.joyliveusb" ]]; then
#               echo "$mnt does not contain a USB key" >&2
#               return 1
#       fi

        umount "$mnt"
}

function usage(){
    echo "usage: $(basename $0) <command> [args]"
    echo ""
    echo "   command: "
    echo "       mount:   mount usbkey"
    echo "          args: 不需要"   
    echo "       unmount: umount usbkey"
    echo "          args: 不需要"  
    echo "       upgrade: 更新 SmartOS pi"
    echo "          args: 可选，pi文件路径，必须为tgz"     
    exit 1
}

function fatal(){
    echo "error: $1"
    exit 2
}

function clean(){    
    echo -n "clean..."
    rm -r ${tmp_path}/test_platform
    rm ${tmp_path}/${gzfile}
    echo " DONE"
}


function upgrade(){        
    mkdir -p ${tmp_path}
    cd ${tmp_path}
    local pi_file=$1
    gzfile=$(basename ${pi_file})
    ext_name=${gzfile##*.}
    [[ "${ext_name}" == "tgz" ]] || fatal "pi 文件 ${gzfile} 不是tgz文件"
    
    if [ "$(echo ${pi_file:0:4}|tr '[:upper:]' '[:lower:]')" == "http" ]; then
        echo -n "下载 ${gzfile}..."
        curl -sO ${pi_file}
        echo " DONE"
        [[ -s "${tmp_path}/${gzfile}" ]] || fatal "pi 文件 不存在: ${pi_file}"       
    else
        [[ -f "${pi_file}" ]] || fatal "pi 文件 不存在: ${pi_file}"
        [[ "${tmp_path}/${gzfile}" == "${pi_file}" ]] || mv ${pi_file} ${tmp_path}/${gzfile}         
    fi
    
    echo -n "解压 pi tgz..."
    mkdir -p ${tmp_path}/test_platform
    rm -rf ${tmp_path}/test_platform/*
    tar xzf ${tmp_path}/${gzfile} -C ${tmp_path}/test_platform/
    echo " DONE"

    current_ver=$(uname -v | cut -d_ -f2)
    echo "当前 pi 版本: ${current_ver}"
    new_ver=$(ls ${tmp_path}/test_platform | grep platform-|cut -f 2 -d"-")
    echo "更新 pi 版本: ${new_ver}"
    
    if [ "${current_ver}" == "${new_ver}" ];then
      echo "pi 版本相同，无需更新"
      clean ${tmp_path}
      exit
    fi
    read -p "确定更新吗? [y/N]:" input_YN
    if [ "${input_YN}" != "y" ]; then
        if [ "${input_YN}" != "Y" ]; then
            clean ${tmp_path}
            exit
        fi
    fi

    mnt_point=/mnt/$(svcprop -p "joyentfs/usb_mountpoint" "svc:/system/filesystem/smartdc:default")
    echo -n "mount usbkey..."
    mount_usb_key ${mnt_point}
    echo " DONE"
    echo -n "复制 pi 文件..."    
    cp ${tmp_path}/test_platform/* ${mnt_point}
    echo " DONE"

    pi_stamp=$(ls ${mnt_point}|grep platform-|cut -f 2 -d"-")

    if [ ${backup_pi} == true ]; then            
        previous=previous
        mkdir -p ${mnt_point}/${previous}
        if [ -d "${mnt_point}/${previous}/platform" ];then
                rm -r ${mnt_point}/${previous}/platform
        fi
        echo -n "备份旧 pi 到 ${previous} && 安装 platform-${pi_stamp}..."
        mv ${mnt_point}/platform ${mnt_point}/${previous}/ && mv ${mnt_point}/platform{-$pi_stamp,}
        echo " DONE"
    else
        echo -n "安装 platform-${pi_stamp}..."
        rm -r ${mnt_point}/platform && mv ${mnt_point}/platform{-$pi_stamp,}
        echo " DONE"
    fi

    clean ${tmp_path}

    echo -n "umount ${usbdrive}..."
    unmount_usb_key ${mnt_point}
    echo " DONE"
    echo "更新成功，重启后生效!"
    # 只有重新启动后才生效
    #reboot
}

# 设置是否备份前一个版本
backup_pi=true
# 临时目录，用来存在pi文件
tmp_path=/var/tmp

# pi文件路径，pi文件必须为tgz，可以为本地路径或远程url（一般只有远程为内网时才设置远程）
# pi_file=https://manta.smartos.top/Joyent_Dev/public/SmartOS/platform-latest.tgz
pi_file=${tmp_path}/platform-latest.tgz

# pi文件路径 默认为已经下载到本地: ${tmp_path}/platform-latest.tgz（推荐）
# 可以通过参数2来修改为直接从远程下载 
[[ "$2" == "" ]] || pi_file=$2

cmd=$1

case $cmd in
    mount) mount_usb_key ;;
    unmount) unmount_usb_key ;;
    upgrade) upgrade "${pi_file}" ;;
    *) usage ;;
esac
